#include "refine/despill3.h"
#include "refine/window_pass.h"
#include <limits>
//#include <iostream>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g3/smatrix.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

typedef unsigned uint;
typedef g3::basic_vector<uint>		vectoru;
typedef g3::basic_smatrix<uint>		smatrixu;
typedef g3::basic_vector<double>	vector;
typedef g3::basic_smatrix<double>	smatrix;
typedef g3::basic_vector<float>		vectorf;

inline vectoru v(byte_color const & c)
	{ return vectoru(c.b, c.g, c.r); }

// reconstructs foreground color from a number of image samples with known alpha
template <typename Raster>
class despiller3
{
public:
	typedef typename Raster::elem_type color;
	typedef typename color::component_type grey_mask_value;
	const static grey_mask_value grey_mask_max_value = 255; //std::numeric_limits<grey_mask_value>::max();	//255
	const static grey_mask_value grey_mask_min_value = 0; //std::numeric_limits<grey_mask_value>::min();	//0
public:
	despiller3(Raster & im_, int whs) : im(im_) {
		clear(); 

		img::bitmap not_max(im.dim());
		for (img::bitmap::iterator i = not_max.begin(); i != not_max.end(); ++i)
			*i = im[i].a != grey_mask_max_value;
		img::bitmap not_max_x;
		img::dilate8(not_max_x, not_max, whs);
		img::combine(not_max_x, not_max, img::bit_and_not());

		img::bitmap not_min(im.dim());
		for (img::bitmap::iterator i = not_min.begin(); i != not_min.end(); ++i)
			*i = im[i].a != grey_mask_min_value;
		img::bitmap not_min_x;
		img::dilate8(not_min_x, not_min, whs);
		img::combine(not_min_x, not_min, img::bit_and_not());

		//modifiable = (not_min & not_max) | not_min_x | not_max_x
		modifiable = not_min;
		img::combine(modifiable, not_max, img::bit_and());
		img::combine(modifiable, not_min_x, img::bit_or());
		img::combine(modifiable, not_max_x, img::bit_or());
	}
	void clear() {
		n = sum_alpha = sum_sqr_alpha = 0;
		sum_I = sum_alpha_I = vectoru::nill;
		sum_IIT.zero();
	}
	size_extent dim() const
		{ return im.dim(); }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p)
		{ add_sample(im[p]); }
	void remove(size_point const & p)
		{ remove_sample(im[p]); }
	void remove_all() 
		{ clear(); }
	void process(size_point const & p, size_rect const &) {
		assert (n > 0);

		{ //check max_condition_number
			double a = sum_sqr_alpha;
			double b = grey_mask_max_value * sum_alpha - sum_sqr_alpha;
			double c = grey_mask_max_value * (grey_mask_max_value * n - 2 * sum_alpha) + sum_sqr_alpha;
			double D = sqr(a - c) + 4 * b * b;
			float max_condition_number = 5;
			if (sqr(max_condition_number + 1) * D > sqr((max_condition_number - 1) * (a + c)))
				return;
		}

		vectorf grad_a = sle0(
			double(n) * smatrix(sum_IIT) - tsquare(vector(sum_I)),
			double(n) * vector(sum_alpha_I) - double(sum_alpha) * vector(sum_I));
		float sqr_grad_a = sqr(grad_a);
		if (sqr_grad_a < std::numeric_limits<float>::epsilon())
			return; //F = I
		float a0 = (sum_alpha - scalar_product(vectorf(sum_I), grad_a)) / n;

		color & I = im[p];
		vectorf F = v(I);
		float a = a0 + scalar_product(F, grad_a);
		//if (a < grey_mask_min_value)
		//	a = grey_mask_min_value;
		//else if (a > grey_mask_max_value)
		//	a = grey_mask_max_value;

		F += (grey_mask_max_value - a) / sqr_grad_a * grad_a; 
		F += vectorf(0.5f, 0.5f, 0.5f); // for proper rounding
		if (F.x < grey_mask_min_value 
		 || F.y < grey_mask_min_value 
		 || F.z < grey_mask_min_value
		 || F.x > grey_mask_max_value 
		 || F.y > grey_mask_max_value 
		 || F.z > grey_mask_max_value)
			return;

		I.b = F.x;
		I.g = F.y;
		I.r = F.z;
	}
private:
	void add_sample(color const & Ic) {
		vectoru I = v(Ic);
		uint alpha = Ic.a;
		++n;
		sum_alpha += alpha;
		sum_sqr_alpha += alpha * alpha;
		sum_I += I;
		sum_alpha_I += alpha * I;
		sum_IIT += tsquare(I);
	}
	void remove_sample(color const & Ic) {
		vectoru I = v(Ic);
		uint alpha = Ic.a;
		--n;
		sum_alpha -= alpha;
		sum_sqr_alpha -= alpha * alpha;
		sum_I -= I;
		sum_alpha_I -= alpha * I;
		sum_IIT -= tsquare(I);
	}
private:
	Raster & im;
	uint n, sum_alpha, sum_sqr_alpha;
	vectoru sum_I, sum_alpha_I;
	smatrixu sum_IIT; ///< sum of (I * I^T)
	img::bitmap modifiable; ///< this pixels may change their alpha theoretically
};

void despill3(byte_image & im, int whs)
{
	despiller3<byte_image> ds(im, whs);
	window_pass(ds, whs);
}

void despill3(argb_byte_image & im, int whs)
{
	despiller3<argb_byte_image> ds(im, whs);
	window_pass(ds, whs);
}

} //namespace refine
